import paramiko
import pandas as pd
import os
from typing import Dict, List, Optional
from contextlib import contextmanager
from logging import basicConfig, INFO, StreamHandler

# 配置日志系统
basicConfig(
    level=INFO,
    format="%(asctime)s - %(levelname)s - %(message)s",
    handlers=[StreamHandler()]
)

class SSHConnectionError(Exception):
    """SSH连接相关异常基类"""
    pass

class CommandExecutionError(Exception):
    """命令执行异常"""
    pass

class FileTransferError(Exception):
    """文件传输异常"""
    pass

class SystemCheckTool:
    def __init__(self, config: Dict):
        self.config = {
            "host": config.get("host", "172.40.0.190"),
            "port": config.get("port", 22),
            "user": config.get("user", "root"),
            "password": config.get("password", "Yu641641."),
            "file_name": config.get("file_name", "system_check.xlsx"),
            "remote_path": config.get("remote_path", "/root/system_check.xlsx"),
            "check_commands": config.get("check_commands", {
                "cpu_model": "cat /proc/cpuinfo | grep 'model name' | head -n 1 | awk -F ':' '{print $2}'",
                "cpu_count": "cat /proc/cpuinfo | grep 'physical id' | uniq | wc -l",
                "cpu_cores": "cat /proc/cpuinfo | grep 'cpu cores' | uniq | awk -F ':' '{print $2}'",
                "disk_usage": "df -h | grep -E 'sd|Filesystem'",
                "memory_info": "free -m | grep -E 'total|Mem'"
            })
        }
        self.ssh_client = None
        self.sftp_client = None

    @contextmanager
    def ssh_connection(self):
        """SSH连接上下文管理器"""
        client = paramiko.SSHClient()
        client.set_missing_host_key_policy(paramiko.AutoAddPolicy())
        try:
            client.connect(
                hostname=self.config["host"],
                port=self.config["port"],
                username=self.config["user"],
                password=self.config["password"]
            )
            self.ssh_client = client
            yield client
        except paramiko.AuthenticationException:
            raise SSHConnectionError("认证失败，请检查用户名和密码") from None
        except paramiko.SSHException as e:
            raise SSHConnectionError(f"SSH连接失败: {str(e)}") from e
        finally:
            if self.ssh_client:
                self.ssh_client.close()

    def _execute_command(self, command: str) -> str:
        """执行命令的通用方法"""
        try:
            stdin, stdout, stderr = self.ssh_client.exec_command(command)
            output = stdout.read().decode('utf-8').strip()
            error = stderr.read().decode('utf-8').strip()
            
            if error:
                raise CommandExecutionError(f"命令执行错误: {error}")
                
            return output
        except Exception as e:
            raise CommandExecutionError(f"命令执行失败: {str(e)}") from e

    def _generate_report(self) -> pd.DataFrame:
        """生成系统检查报告"""
        results = {}
        for key, cmd in self.config["check_commands"].items():
            try:
                output = self._execute_command(cmd)
                results[key] = self._process_output(key, output)
            except CommandExecutionError as e:
                results[key] = f"获取失败: {str(e)}"
        
        return pd.DataFrame([results])

    def _process_output(self, command: str, raw_output: str) -> str:
        """处理不同命令的输出格式"""
        processors = {
            "cpu_model": lambda x: x.split(':')[1].strip(),
            "cpu_count": int,
            "cpu_cores": int,
            "disk_usage": lambda x: "\n".join(x.splitlines()[1:]),
            "memory_info": lambda x: "\n".join(x.splitlines()[1:])
        }
        return processors.get(command, str)(raw_output)

    def _upload_file(self) -> None:
        """上传报告文件到远程服务器"""
        try:
            with self.ssh_connection() as client:
                transport = client.get_transport()
                sftp = paramiko.SFTPClient.from_transport(transport)
                
                remote_dir = os.path.dirname(self.config["remote_path"])
                sftp.mkdir(remote_dir, exist_ok=True)
                
                sftp.put(self.config["file_name"], self.config["remote_path"])
                sftp.chmod(self.config["remote_path"], 0o644)
        except FileTransferError as e:
            raise FileTransferError(f"文件上传失败: {str(e)}") from e

import logging

# 初始化 Logger
LOGGER = logging.getLogger(__name__)

class YourClass:
    def upload_file(self):
        # ... 你的上传逻辑 ...
        LOGGER.info(f"文件上传成功至：{self.config['remote_path']}")
    def generate_and_upload_report(self) -> None:
        """完整执行流程"""
        try:
            # 生成报告
            report_df = self._generate_report()
            report_path = os.path.abspath(self.config["file_name"])
            report_df.to_excel(report_path, index=False)
            
            LOGGER.info(f"报告生成成功: {report_path}")

            # 上传报告
            self._upload_file()
            LOGGER.info(f"文件上传成功至: {self.config['remote_path']}")

        except SSHConnectionError as e:
            LOGGER.error(f"连接错误: {str(e)}")
        except FileTransferError as e:
            LOGGER.error(f"传输错误: {str(e)}")
        except Exception as e:
            LOGGER.error(f"未知错误: {str(e)}")

if __name__ == "__main__":
    config = {
        "host": "172.40.0.190",
        "password": "Yu641641.",
        "file_name": "server_check_report.xlsx",
        "remote_path": "/var/log/server_checks/",
        "check_commands": {
            "os_version": "cat /etc/os-release | grep PRETTY_NAME",
            "kernel_info": "uname -a",
            "memory_total": "free -m | grep Mem | awk '{print $2}'",
            "disk_space": "df -h --total | grep total"
        }
    }

    tool = SystemCheckTool(config)
    tool.generate_and_upload_report()